import torch
import torch.nn.functional as F
import torch_geometric.transforms as T

from GCN import GCN
from util import load_dataset

name = 'Cora'
hidden = 32
lr = 0.01
weight_decay = 5e-4

dataset = load_dataset(name)
dataset.transform = T.NormalizeFeatures()

model = GCN(dataset.num_node_features, hidden, dataset.num_classes)
data = dataset[0]

if torch.cuda.is_available():
    model = model.cuda()
    data = data.cuda()

optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

best_val_loss = 1000
early_stop = 200
patience = 200
best_model_state = None

for epoch in range(2000):
    model.train()
    optimizer.zero_grad()
    output = model(data)
    loss = F.nll_loss(output[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

    model.eval()
    output = model(data)
    val_loss = F.nll_loss(output[data.val_mask], data.y[data.val_mask])
    print(f"epoch:{epoch}   Val_loss:{val_loss:.4f}")
    if best_val_loss > val_loss:
        best_val_loss = val_loss
        patience = early_stop
        best_model_state = model.state_dict()
    else:
        patience -= 1
    if epoch > early_stop and patience <= 0:
        break

model.load_state_dict(best_model_state)
model.eval()
_, pred = model(data).max(dim=1)
correct = pred[data.test_mask].eq(data.y[data.test_mask]).sum()
acc = correct / int(data.test_mask.sum())
print("accuracy= {:.4f}".format(acc))
